package tcp

import (
	"context"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"tunnel_manager/core"

	"github.com/quic-go/quic-go"
	"golang.org/x/sync/semaphore"
)

// ========================================
// TCP协议处理器实现
// 实现 core.ProtocolHandler 接口
// ========================================

// Handler TCP协议处理器
type Handler struct {
	tunnelID   string
	targetHost string
	targetPort string

	// 并发控制
	concurrencySem *semaphore.Weighted
	maxConcurrency int64
	activeTasks    int64

	// 连接池
	connPool     sync.Map
	connPoolSize int32
	maxPoolSize  int32

	// 统计信息
	stats core.ProtocolStats
}

// NewHandler 创建TCP协议处理器
func NewHandler(tunnelID, targetHost, targetPort string) *Handler {
	maxConcurrency := int64(100)

	handler := &Handler{
		tunnelID:       tunnelID,
		targetHost:     targetHost,
		targetPort:     targetPort,
		concurrencySem: semaphore.NewWeighted(maxConcurrency),
		maxConcurrency: maxConcurrency,
		activeTasks:    0,
		maxPoolSize:    50,
	}

	log.Printf("✅ [TCP协议] 处理器已创建: Target=%s:%s, MaxConcurrency=%d",
		targetHost, targetPort, maxConcurrency)

	return handler
}

// HandleRequest 实现 ProtocolHandler 接口
func (h *Handler) HandleRequest(stream quic.Stream, msg *core.TunnelMessage) error {
	requestID := msg.RequestID
	log.Printf("📥 [TCP协议] 收到请求: RequestID=%s, Size=%d bytes",
		requestID[:8]+"...", len(msg.Body))

	// 并发控制
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	if err := h.concurrencySem.Acquire(ctx, 1); err != nil {
		log.Printf("❌ [TCP协议] 并发超限: %v", err)
		errorMsg := core.NewErrorMessage(h.tunnelID, requestID, 1, "并发超限")
		errorMsg.Write(stream)
		return err
	}
	defer h.concurrencySem.Release(1)

	atomic.AddInt64(&h.activeTasks, 1)
	defer atomic.AddInt64(&h.activeTasks, -1)
	atomic.AddInt64(&h.stats.TotalRequests, 1)

	// 建立到目标服务的TCP连接
	targetConn, err := h.getOrCreateConnection(ctx)
	if err != nil {
		log.Printf("❌ [TCP协议] 连接目标失败: %v", err)
		atomic.AddInt64(&h.stats.ErrorCount, 1)
		errorMsg := core.NewErrorMessage(h.tunnelID, requestID, 2, fmt.Sprintf("连接失败: %v", err))
		errorMsg.Write(stream)
		return err
	}

	// 转发请求数据
	if len(msg.Body) > 0 {
		_, err = targetConn.Write(msg.Body)
		if err != nil {
			log.Printf("❌ [TCP协议] 写入目标失败: %v", err)
			h.removeConnection(targetConn)
			targetConn.Close()
			atomic.AddInt64(&h.stats.ErrorCount, 1)
			errorMsg := core.NewErrorMessage(h.tunnelID, requestID, 3, fmt.Sprintf("写入失败: %v", err))
			errorMsg.Write(stream)
			return err
		}
		atomic.AddInt64(&h.stats.TotalBytesRecv, int64(len(msg.Body)))
	}

	// 双向流式转发
	return h.bidirectionalForward(ctx, stream, targetConn, requestID)
}

// getOrCreateConnection 获取或创建TCP连接
func (h *Handler) getOrCreateConnection(ctx context.Context) (net.Conn, error) {
	targetAddr := net.JoinHostPort(h.targetHost, h.targetPort)

	dialer := &net.Dialer{
		Timeout:   10 * time.Second,
		KeepAlive: 30 * time.Second,
	}

	conn, err := dialer.DialContext(ctx, "tcp", targetAddr)
	if err != nil {
		return nil, fmt.Errorf("拨号失败: %w", err)
	}

	if tcpConn, ok := conn.(*net.TCPConn); ok {
		tcpConn.SetKeepAlive(true)
		tcpConn.SetKeepAlivePeriod(30 * time.Second)
		tcpConn.SetNoDelay(true)
	}

	currentSize := atomic.LoadInt32(&h.connPoolSize)
	if currentSize < h.maxPoolSize {
		connKey := fmt.Sprintf("%p", conn)
		h.connPool.Store(connKey, conn)
		atomic.AddInt32(&h.connPoolSize, 1)
	}

	log.Printf("🔗 [TCP协议] 建立新连接: %s", targetAddr)
	return conn, nil
}

// removeConnection 从连接池移除连接
func (h *Handler) removeConnection(conn net.Conn) {
	connKey := fmt.Sprintf("%p", conn)
	if _, loaded := h.connPool.LoadAndDelete(connKey); loaded {
		atomic.AddInt32(&h.connPoolSize, -1)
	}
}

// bidirectionalForward 双向流式转发
func (h *Handler) bidirectionalForward(
	ctx context.Context,
	stream quic.Stream,
	targetConn net.Conn,
	requestID string,
) error {
	var wg sync.WaitGroup
	errChan := make(chan error, 2)

	// 从目标读取 → 发送到Stream
	wg.Add(1)
	go func() {
		defer wg.Done()
		err := h.forwardTargetToStream(ctx, stream, targetConn, requestID)
		if err != nil && err != io.EOF {
			log.Printf("⚠️  [TCP协议] 目标→Stream转发结束: %v", err)
		}
		errChan <- err
	}()

	// 从Stream读取 → 发送到目标
	wg.Add(1)
	go func() {
		defer wg.Done()
		err := h.forwardStreamToTarget(ctx, stream, targetConn, requestID)
		if err != nil && err != io.EOF {
			log.Printf("⚠️  [TCP协议] Stream→目标转发结束: %v", err)
		}
		errChan <- err
	}()

	go func() {
		wg.Wait()
		close(errChan)
	}()

	var lastErr error
	for err := range errChan {
		if err != nil && err != io.EOF {
			lastErr = err
		}
	}

	targetConn.Close()
	h.removeConnection(targetConn)

	if lastErr == nil {
		atomic.AddInt64(&h.stats.SuccessCount, 1)
	} else {
		atomic.AddInt64(&h.stats.ErrorCount, 1)
	}

	return lastErr
}

// forwardTargetToStream 从目标服务读取数据并发送到Stream
func (h *Handler) forwardTargetToStream(
	ctx context.Context,
	stream quic.Stream,
	targetConn net.Conn,
	requestID string,
) error {
	buffer := make([]byte, 32*1024)
	totalSent := int64(0)

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		targetConn.SetReadDeadline(time.Now().Add(30 * time.Second))
		n, err := targetConn.Read(buffer)
		if n > 0 {
			respMsg := &core.TunnelMessage{
				Version:      core.ProtocolVersion,
				MessageType:  core.MessageTypeResponse,
				ProtocolType: core.ProtocolTCP,
				TunnelID:     h.tunnelID,
				RequestID:    requestID,
				Body:         buffer[:n],
				BodyLength:   uint64(n),
			}

			if err := respMsg.Write(stream); err != nil {
				log.Printf("❌ [TCP协议] 写入Stream失败: %v", err)
				return err
			}

			totalSent += int64(n)
			atomic.AddInt64(&h.stats.TotalBytesSent, int64(n))
		}

		if err != nil {
			if err == io.EOF {
				log.Printf("✅ [TCP协议] 目标连接正常关闭, 总计: %.2f KB",
					float64(totalSent)/1024)
				return io.EOF
			}
			return fmt.Errorf("读取目标失败: %w", err)
		}
	}
}

// forwardStreamToTarget 从Stream读取数据并发送到目标服务
func (h *Handler) forwardStreamToTarget(
	ctx context.Context,
	stream quic.Stream,
	targetConn net.Conn,
	requestID string,
) error {
	totalReceived := int64(0)

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		msg, err := core.ReadMessage(&stream)
		if err != nil {
			if err == io.EOF {
				return io.EOF
			}
			return fmt.Errorf("读取Stream失败: %w", err)
		}

		if msg.MessageType != core.MessageTypeRequest && msg.MessageType != core.MessageTypeResponse {
			continue
		}

		if len(msg.Body) > 0 {
			targetConn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			_, err = targetConn.Write(msg.Body)
			if err != nil {
				return fmt.Errorf("写入目标失败: %w", err)
			}

			totalReceived += int64(len(msg.Body))
			atomic.AddInt64(&h.stats.TotalBytesRecv, int64(len(msg.Body)))
		}
	}
}

// GetProtocolType 实现 ProtocolHandler 接口
func (h *Handler) GetProtocolType() core.ProtocolType {
	return core.ProtocolTCP
}

// Close 实现 ProtocolHandler 接口
func (h *Handler) Close() error {
	h.connPool.Range(func(key, value interface{}) bool {
		if conn, ok := value.(net.Conn); ok {
			conn.Close()
		}
		h.connPool.Delete(key)
		return true
	})
	atomic.StoreInt32(&h.connPoolSize, 0)
	log.Printf("🛑 [TCP协议] 处理器已关闭")
	return nil
}

// GetStats 获取统计信息
func (h *Handler) GetStats() core.ProtocolStats {
	return core.ProtocolStats{
		TotalRequests:  atomic.LoadInt64(&h.stats.TotalRequests),
		ActiveRequests: atomic.LoadInt64(&h.activeTasks),
		TotalBytesSent: atomic.LoadInt64(&h.stats.TotalBytesSent),
		TotalBytesRecv: atomic.LoadInt64(&h.stats.TotalBytesRecv),
		SuccessCount:   atomic.LoadInt64(&h.stats.SuccessCount),
		ErrorCount:     atomic.LoadInt64(&h.stats.ErrorCount),
	}
}
